package com.tuniu.agents.common.logger;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.TraceLogVo;
import com.tuniu.agents.common.service.ChatTraceLogService;
import com.tuniu.agents.message.TraceLogger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

@Service
public class TraceDateLogger implements TraceLogger, InitializingBean, DisposableBean {

    private static final Logger logger = LoggerFactory.getLogger(TraceDateLogger.class);
    private static final BlockingQueue<TraceLogVo> queue = new LinkedBlockingQueue<>();
    private static final ExecutorService executorService = Executors.newSingleThreadExecutor();
    private final AtomicBoolean shuttingDown = new AtomicBoolean(false);

    @Autowired
    private ChatTraceLogService chatTraceLogService;

    @Autowired
    private ChatEnvironmentMemory chatEnvironmentMemory;

    @Override
    public void afterPropertiesSet() throws Exception {
        // Create a single worker thread to process the queue
        executorService.submit(() -> {
            while (!Thread.currentThread().isInterrupted() && !shuttingDown.get()) {
                try {
                    // Use short timeout polling instead of infinite blocking
                    TraceLogVo traceLogVo = queue.poll(100, TimeUnit.MILLISECONDS);
                    if (traceLogVo != null) {
                        switch (traceLogVo.getOperationType()){
                            case "chatTraceLogSaveTrace":
                                saveTraceAsync(traceLogVo);
                                break;
                            case "chatTraceLogUpdateRecentLLMTrace":
                                updateRecentLLMTraceSync(traceLogVo);
                                break;
                        }
                    }
                } catch (InterruptedException e) {
                    // Normal interrupt handling, set interrupt flag and exit
                    Thread.currentThread().interrupt();
                    logger.info("Worker thread interrupted, exiting");
                    break;
                } catch (Exception e) {
                    logger.error("Error processing trace log", e);
                }
            }
            
            // Process remaining logs in queue when shutdown flag is set
            if (shuttingDown.get()) {
                try {
                    TraceLogVo traceLogVo;
                    logger.info("Processing remaining logs in queue before worker thread exits");
                    // Process all remaining logs in queue non-blockingly
                    while ((traceLogVo = queue.poll()) != null) {
                        switch (traceLogVo.getOperationType()){
                            case "chatTraceLogSaveTrace":
                                saveTraceAsync(traceLogVo);
                                break;
                            case "chatTraceLogUpdateRecentLLMTrace":
                                updateRecentLLMTraceSync(traceLogVo);
                                break;
                        }
                    }
                    logger.info("All remaining logs processed by worker thread");
                } catch (Exception e) {
                    logger.error("Error processing remaining logs", e);
                }
            }
            
            logger.info("Worker thread exiting");
        });
        logger.info("TraceDateLogger initialized and started");
    }

    @Override
    public void destroy() throws Exception {
        // Mark as shutting down
        logger.info("TraceDateLogger shutdown initiated");
        shuttingDown.set(true);
        
        // Wait for the queue to be empty, indicating all logs have been processed
        try {
            logger.info("Waiting for log queue to be empty");
            
            // Poll to check if queue is empty, wait up to 60 seconds
            long startTime = System.currentTimeMillis();
            while (!queue.isEmpty() && System.currentTimeMillis() - startTime < 60000) {
                Thread.sleep(100);
            }
            
            // Check final status
            logger.info("Log queue is empty, all logs have been processed");
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            logger.warn("Interrupted while waiting for log queue to empty");
        }
        
        // Now it's safe to shut down the thread pool
        executorService.shutdown();
        
        try {
            // Give thread pool a reasonable time to complete tasks
            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                logger.info("Worker thread did not terminate in time, forcing shutdown");
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            logger.info("Shutdown interrupted, forcing immediate shutdown");
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
        
        logger.info("TraceDateLogger shutdown complete");
    }

    @Override
    public void saveTrace(TraceLogVo traceLogVo) {
        // If shutting down, don't accept new logs
        if (shuttingDown.get()) {
            logger.info("TraceDateLogger is shutting down, rejecting new log");
            return;
        }
        
        try {
            String conversationId = traceLogVo.getConversationId();
            Object countObject =  chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.PARAM_SPAN_ID_COUNT);
            Integer count = countObject != null ? Integer.parseInt(countObject.toString()) : 0;
            String spanId = count + "-" + traceLogVo.getAgentId();
            Object parentSpanId = chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.PARAM_PARENT_SPAN_ID);
            chatEnvironmentMemory.add(conversationId, Map.of(AgentsConstants.AgentsConversationEnvironmentKey.PARAM_SPAN_ID_COUNT, count + 1,
                    AgentsConstants.AgentsConversationEnvironmentKey.PARAM_PARENT_SPAN_ID, spanId));
            traceLogVo.setSpanId(spanId);
            traceLogVo.setParentId(parentSpanId != null ? parentSpanId.toString() : "");
            
            // Check shutdown status again, use non-blocking offer to avoid blocking during shutdown
            if (!shuttingDown.get()) {
                if (!queue.offer(traceLogVo)) {
                    logger.warn("Failed to add trace log to queue - queue full");
                }
            }
        } catch (Exception e) {
            logger.error("Failed to process trace log", e);
        }
    }

    private void saveTraceAsync(TraceLogVo traceLogVo) {
        try {
            chatTraceLogService.saveTrace(traceLogVo);
        } catch (Exception e) {
            logger.error("Failed to save trace log asynchronously", e);
        }
    }

    @Override
    public void updateRecentLLMTrace(TraceLogVo traceLogVo) {
        if (shuttingDown.get()) {
            logger.info("TraceDateLogger is shutting down, rejecting new log");
            return;
        }
        try {
            if (!shuttingDown.get()) {
                if (!queue.offer(traceLogVo)) {
                    logger.warn("Failed to add trace log to queue - queue full");
                }
            }
        } catch (Exception e) {
            logger.error("Failed to process trace log", e);
        }
    }

    public void updateRecentLLMTraceSync(TraceLogVo traceLogVo) {
        try {
            chatTraceLogService.updateRecentLLMTrace(traceLogVo);
        } catch (Exception e) {
            logger.error("Failed to update trace log asynchronously", e);
        }
    }
}